///////////////////////////////////////////////////////////////////////////////
// Name:        objectdet.cpp
// Purpose:     Main function for object detection
// Created:     2006
// Copyright:   (c) 2006, Ivan Laptev, TT Solutions
// License:     GPL, Please contact us for commercial license.
//              http://www.tt-solutions.com/
///////////////////////////////////////////////////////////////////////////////

#include "objectdet/common.h"

#include "objectdet/adacascadeclass.h"
#include "objectdet/data.h"
#include "objectdet/detectobjects.h"
#include "objectdet/dump.h"
#include "objectdet/integralhistogram.h"
#include "objectdet/localhistintegral.h"
#include "objectdet/prune_subwindows.h"
#include "objectdet/profiling.h"
#include "objectdet/xmlloadmodel.h"

void print_usage_info()
{
    std::cout
        << "Usage: objectdet\n"
        << "    -s          : streaming input from camera\n"
        << "    -i filename : image input from file\n"
        << "    -c filename : classifier file name (xml)\n"
        << "    -o filename : output xml file if output to disk is required\n"
        << "    -n integer  : number of pyramid levels per octave (=10 default)\n"
        << "    -r integer  : pyramid level step for image processing (=1 default)\n"
        << "    -e double   : exponent value to define image scale at the finest\n"
        << "                  pyramid level: scale_factor = 1/2^exp (=0 default)\n"
        << "    -q          : quite flag disables graphical/textual output\n"
        << "    -h          : this message\n";
}

int do_main(int argc, char* argv[])
{
    // using namespace std;
    std::string imgfname   = "images/000537.png";
    std::string classfname = "classifiers/cabmodel_interm_nst40_VOC06car01train5_trainval_orienthistmod.xml";
    std::string resfname   = "results/detrect.xml";
    size_t nlevelsinoctave = 10;
    size_t recomplevelstep = 1;
    float expvalinit     = 0.0;
    bool streamingflag   = false;
    bool saveresflag     = false;
    bool quietflag       = false;

    /**
     *  Process command line arguments
     */
    for(int i = 1; i < argc; ++i)  // Skip argv[0] (program name).
    {
        if (strcmp(argv[i], "-s") == 0)  // camera input
        {
            streamingflag = true;
        }
        else if (strcmp(argv[i], "-i") == 0)  // image input
        {
            if (++i < argc)
            {
                imgfname = argv[i];
            }
        }
        else if (strcmp(argv[i], "-c") == 0)  // classifier file name
        {
            if (++i < argc)
            {
                classfname = argv[i];
            }
        }
        else if (strcmp(argv[i], "-o") == 0)
        {  // file name for xml result output
            if (++i < argc)
            {
                saveresflag = true;
                resfname = argv[i];
            }
        }
        else if (strcmp(argv[i], "-n") == 0)
        {  // num. of pyramid levels in octave
            if (++i < argc)
            {
                nlevelsinoctave = boost::lexical_cast<int>(argv[i]);
            }
        }
        else if (strcmp(argv[i], "-r") == 0)
        {  // level-step for image processing
            if (++i < argc)
            {
                recomplevelstep = boost::lexical_cast<int>(argv[i]);
            }
        }
        else if (strcmp(argv[i], "-e") == 0)
        {  // exponent value for the 0 pyramid level
            if (++i < argc)
            {
                expvalinit = boost::lexical_cast<float>(argv[i]);
            }
        }
        else if (strcmp(argv[i], "-q") == 0)
        { // quiet flag
            quietflag = true;
        }
        else if (strcmp(argv[i], "-h") == 0)
        { // quiet flag
            print_usage_info();
            exit(0);
        }
        else
        {
            print_usage_info();
            exit(1);
        }
    }

    if (!quietflag)
    {
        std::cout << "Start object detection using parameters:\n";
        if (streamingflag)
        {
            std::cout << "  Camera input\n";
        }
        else
        {
            std::cout << "  Image file iput:       " << imgfname << "\n";
        }
        std::cout << "  Classifier file name:  " << classfname << "\n";
        if (saveresflag)
        {
            std::cout << "  Output in file:        " << resfname << "\n";
        }
    }

    //*** initialize camera input
    CvCapture* capture = 0;
    pImage src;
    int videodev=1;
    if (streamingflag)
    {
        capture = cvCaptureFromCAM(videodev);    
        if(!capture)
        {
            std::ostringstream oss;
            oss << "cannot initialize video device '" << videodev << "'";
            throw std::runtime_error(oss.str());
        }
        cvGrabFrame(capture);
        pImage frame = dummyGuard(cvRetrieveFrame(capture));
        if(!frame)
        {
            std::ostringstream oss;
            oss << "cannot read frame from video device '" << videodev << "'";
            throw std::runtime_error(oss.str());
        }
        src = guard(cvCreateImage(cvSize(frame->width,frame->height), IPL_DEPTH_8U, frame->nChannels));
        if( frame->origin == IPL_ORIGIN_TL )
        {
            CopyMatrix(ImageAsMatrix(frame), ImageAsMatrix(src));
        }
        else
        {
            FlipMatrix(ImageAsMatrix(frame), ImageAsMatrix(src), 0);
        }
    }    
    //*** initialise image file input
    else
    {
        src = guard(cvLoadImage(imgfname.c_str(), 0));
        if (!src)
        {
            std::ostringstream oss;
            oss << "Could not read image " << imgfname << "\n";
            throw std::runtime_error(oss.str());
        }
    }
    //*** load model
    Model model;
    xmlloadmodel(classfname,model);
    Tick("Model loading");
    //printf("Load %d-cascade classifier from %s\n",model.size(),classfname);

    //*** set image processing options
    PixClass pclass;
    pMat cmeans;
    pMat covmat;
    pMat cmeansnorm;
    if (model.begin()->histtype == "orienthistmod")
    {
        //*** use histograms of gradient orientation
        pclass = TEXTONS_SC12_GDORDER2;
    }
    else
    {
        //*** WARNING: no other histogram types are supported in this version
        std::ostringstream oss;
        oss << "Histograms of type '"
            << model.begin()->histtype
            << "' are not supported by the current version\n";
        throw std::runtime_error(oss.str());
        //*** read textons
        //std::string tfname = "textons/textons_numberplate_set1_tnum12_scn2_ftrn5.xml";
        //pFileStorage tfsr = OpenFileStorage(tfname, RWM_Read);
        //pMat cmeans = ReadMatrixFromFile(tfsr, "cmeans");
        //pMat covmat = ReadMatrixFromFile(tfsr, "sqrtinvcovmat");
        //pMat cmeansnorm = ReadMatrixFromFile(tfsr, "cmeansnorm");
    }

    //*** set image scales 
    ScaleFctArray scarray;
    if (nlevelsinoctave<=0)
        nlevelsinoctave=5;
    if (recomplevelstep<1)
        recomplevelstep=1;
    for(double expval = expvalinit;
        expval <= 10.0;
        expval += (double)recomplevelstep*1.0/(double)nlevelsinoctave)
    {
        scarray.push_back(1./pow(2.,expval));
    }

    //*** set histogram blocks
    HistBlockSizeArray bsarray;
    double bszx=(double)(model.begin()->histblock_size_x);
    double bszy=(double)(model.begin()->histblock_size_y);
	for (std::size_t i = 0; i < recomplevelstep; ++i)
    {
        double expval = (double)i/(double)nlevelsinoctave;
        double fct = pow(2.0,expval);
        bsarray.push_back(HistBlockSize(fct*bszx, fct*bszy));
        //*** add robustness to affine transformations
        //bsarray.push_back(HistBlockSize(1.2*fct*bszx, fct*bszy));
        //bsarray.push_back(HistBlockSize(fct*bszx, 1.2*fct*bszy));
    }
    
    CvFont* font = new CvFont();
    CvScalar col;
    cvInitFont(font, CV_FONT_HERSHEY_PLAIN, 1., 1.);
    pImage imgdisp = guard(cvCreateImage(cvSize(src->width, src->height), 8, 3));

    SubWindows subwindows;

    for(;;)
    {
        if (streamingflag)
        {
            cvGrabFrame(capture);
            pImage frame = dummyGuard(cvRetrieveFrame(capture));
            if(!frame)
            {
                std::ostringstream oss;
                oss << "cannot read frame from video device '" << videodev << "'";
                throw std::runtime_error(oss.str());
            }
            if( frame->origin == IPL_ORIGIN_TL )
            {
                CopyMatrix(ImageAsMatrix(frame), ImageAsMatrix(src));
            }
            else
            {
                FlipMatrix(ImageAsMatrix(frame), ImageAsMatrix(src), 0);
            }
        }

        //*** detect objects
        subwindows = DetectObjects(src,
                                   model,
                                   cmeans,
                                   covmat,
                                   cmeansnorm,
                                   pclass,
                                   bsarray,
                                   scarray);

        //*** copy detection rectangles
        int rnum = (int)subwindows.size();
        pMat detrect = CreateMat(rnum > 0 ? rnum : 1, 5, CV_32FC1);
        if (rnum>0)
        {
            int wcount=0;
            for (SubWindows::iterator wit = subwindows.begin(); wit != subwindows.end(); ++wit)
            {
                SetReal2D(detrect, wcount, 0, wit->x);
                SetReal2D(detrect, wcount, 1, wit->y);
                SetReal2D(detrect, wcount, 2, wit->w);
                SetReal2D(detrect, wcount, 3, wit->h);
                SetReal2D(detrect, wcount, 4, wit->cluster_size);
                wcount++;
            }
        } else {
            SetReal2D(detrect, 0, 4, -1);
        }

        //*** display detection results
        if (!quietflag)
        {
            if (src->nChannels == 1)
            {
                pMat src_mat = ImageAsMatrix(src);
                Merge(src_mat, src_mat, src_mat, pMat(),
                      ImageAsMatrix(imgdisp));
            }
            else
            {
                CopyMatrix(ImageAsMatrix(src), ImageAsMatrix(imgdisp));
            }

            if (rnum>0)
            {
                for (int i=0; i<rnum; i++)
                {
                    double rx = GetReal2D(detrect, i, 0);
                    double ry = GetReal2D(detrect, i, 1);
                    double rw = GetReal2D(detrect, i, 2);
                    double rh = GetReal2D(detrect, i, 3);
                    double rc = GetReal2D(detrect, i, 4);

                    col = (rc<5) ? cvScalar(100, 100, 0) : cvScalar(255,255,0);

                    if (rc>=5)
                    {
                        cvRectangle(imgdisp.get(), cvPoint((int)rx,(int)ry),cvPoint((int)(rx+rw),(int)(ry+rh)),col,3);
                        std::ostringstream confstr;
                        confstr << "conf: " << (int)rc;
                        cvPutText(imgdisp.get(), confstr.str().c_str(), cvPoint((int)rx, (int)ry), font, cvScalar(0,0,255));
                    }
                }
            }
            cvNamedWindow("Results", 1);
            cvShowImage("Results", imgdisp.get());
        }

        if (saveresflag)
        {
            if (!quietflag)
                std::cout << "Saving detection rectangles to " << resfname << "\n";
            pFileStorage fsw = OpenFileStorage(resfname, RWM_Write);
            if (!fsw)
            {
                std::ostringstream oss;
                oss << "Could not write results to " << resfname << "\n";
                throw std::runtime_error(oss.str());
            }
            WriteMatrixToFile(fsw, "detrect", detrect);
        }

        if (cvWaitKey(10)>0) break;
        
        if (!streamingflag)
        {
            if (!quietflag)
                cvWaitKey(0);
            break;
        }
    }

    return 0;
}

int main(int argc, char* argv[])
{
    try
    {
        return do_main(argc, argv);
    }
    catch(std::exception const & e)
    {
        std::cerr << e.what();
    }
    return 1;
}

